# src/Dataset/dataset_utils_ms.py
# 这是从 PyTorch 版本迁移过来的辅助函数库

import numpy as np
import torch  # [注意] 我们只在这里临时使用torch.load来读取旧的.pt统计文件
import os
from tqdm import tqdm

# --- 完全复用的函数 (与框架无关) ---

one_letter = {
    'VAL':'V', 'ILE':'I', 'LEU':'L', 'GLU':'E', 'GLN':'Q',
    'ASP':'D', 'ASN':'N', 'HIS':'H', 'TRP':'W', 'PHE':'F', 'TYR':'Y',
    'ARG':'R', 'LYS':'K', 'SER':'S', 'THR':'T', 'MET':'M', 'ALA':'A',
    'GLY':'G', 'PRO':'P', 'CYS':'C'
}

allowable_amino_acids = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE', 'LEU', 'LYS', 'MET',
                         'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL', 'HIP', 'HIE', 'TPO', 'HID', 'LEV',
                         'MEU', 'PTR', 'GLV', 'CYT', 'SEP', 'HIZ', 'CYM', 'GLM', 'ASQ', 'TYS', 'CYX', 'GLZ', 'misc']

def log(*args, **kwargs):
    print(*args, **kwargs)

def dihedral(p0, p1, p2, p3):
    p0, p1, p2, p3 = np.array(p0), np.array(p1), np.array(p2), np.array(p3)
    b0, b1, b2 = -1.0*(p1 - p0), p2 - p1, p3 - p2
    norm_b1 = np.linalg.norm(b1)
    if norm_b1 < 1e-7: return 0.0 # 避免除零
    b1 /= norm_b1
    v = b0 - np.dot(b0, b1)*b1
    w = b2 - np.dot(b2, b1)*b1
    return np.degrees(np.arctan2(np.dot(np.cross(b1, v), w), np.dot(v, w)))

def safe_index(l, e):
    # 完全复用原始逻辑
    try:
        return l.index(e)
    except:
        return len(l) - 1 # 'misc'

# --- 需要从 PyTorch 迁移的函数 ---

def one_hot_res(residx, num_residue_type=20):
    """
    [迁移] 忠实迁移原始 one_hot_res 的逻辑。
    如果残基索引超出了标准范围，则返回None，表示处理失败。
    返回 NumPy 数组而不是 Torch Tensor。
    """
    if residx >= num_residue_type:
        return None  # 这是一个非标准氨基酸，返回失败信号
    
    one_hot = np.zeros(num_residue_type, dtype=np.float32)
    one_hot[residx] = 1.0
    return one_hot

class NormalizeProtein:
    """
    [迁移] NormalizeProtein 类的 MindSpore/NumPy 版本。
    它现在操作的是包含 NumPy 数组的字典。
    """
    def __init__(self, filename=None, mean=None, std=None):
        if filename:
            # 假设统计文件是torch保存的，我们只用它来加载数据
            # 更好的长期方案是您用NumPy重新保存这个文件为 .npy 格式
            stats = torch.load(filename)
            self.mean = stats['mean'].numpy()
            self.std = stats['std'].numpy()
        elif mean is not None and std is not None:
            self.mean = mean
            self.std = std
        else:
            raise ValueError("Either filename or mean/std must be provided.")

    def __call__(self, graph_dict):
        """
        操作字典而不是PyG的Data对象
        """
        # 假设 x 的前20列是one-hot，不需要归一化
        # 第20列是sasa, 第21列是bfactor, 之后是dihedral
        # 这里的索引需要根据您的特征拼接顺序严格对应
        # 从原始代码看，dihedral特征似乎没有被归一化
        
        # 获取需要归一化的特征切片 (SASA 和 B-Factor)
        # 假设它们在第20和21列
        features_to_norm = graph_dict['x'][:, 20:22]
        
        # 从加载的统计数据中获取对应的均值和标准差
        # 原始get_stat可能保存了所有特征的统计，这里只取需要的
        mean_slice = self.mean[20:22]
        std_slice = self.std[20:22]

        graph_dict['x'][:, 20:22] = (features_to_norm - mean_slice) / (std_slice + 1e-9)
        return graph_dict